package com.yeskery.nut.extend.mcp.server;

import com.yeskery.nut.core.MediaType;
import com.yeskery.nut.extend.mcp.core.Prompt;
import com.yeskery.nut.extend.mcp.core.Resource;
import com.yeskery.nut.extend.mcp.core.Tool;
import com.yeskery.nut.extend.mcp.core.schema.ResourceSchema;

import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Mcp服务能力
 * @author sprout
 * 2025-04-18 11:18
 */
public class McpServerCapability {

    /** 工具元数据集合 */
    private final Map<String, McpToolMetadata> toolsMap = new HashMap<>();

    /** 资源uri集合，防止重复uri添加 */
    private final Set<String> resoureUriSet = new HashSet<>();

    /** 资源元数据集合 */
    private final List<McpResourceMetadata> resources = new ArrayList<>();

    /** 提示元数据集合 */
    private final Map<String, McpPromptMetadata> promptsMap = new HashMap<>();

    /**
     * 获取工具集合
     * @return 工具集合
     */
    public List<Tool> getTools() {
        return Collections.unmodifiableList(toolsMap.entrySet()
                .stream()
                .sorted(Map.Entry.comparingByKey())
                .map(Map.Entry::getValue)
                .map(McpToolMetadata::getTool)
                .collect(Collectors.toList()));
    }

    /**
     * 添加工具元数据
     * @param toolMetadataCollection 工具元数据集合
     */
    public void addToolMetadata(McpToolMetadata... toolMetadataCollection) {
        for (McpToolMetadata mcpToolMetadata : toolMetadataCollection) {
            toolsMap.put(mcpToolMetadata.getTool().getTool().getName(), mcpToolMetadata);
        }
    }

    /**
     * 获取工具元数据
     * @param name 工具名称
     * @return 工具元数据
     */
    public McpToolMetadata getToolMetadata(String name) {
        return toolsMap.get(name);
    }

    /**
     * 获取提示集合
     * @return 提示集合
     */
    public List<Prompt> getPrompts() {
        return Collections.unmodifiableList(promptsMap.entrySet()
                .stream()
                .sorted(Map.Entry.comparingByKey())
                .map(Map.Entry::getValue)
                .map(McpPromptMetadata::getPrompt)
                .collect(Collectors.toList()));
    }

    /**
     * 添加提示元数据
     * @param promptMetadataCollection 提示元数据集合
     */
    public void addPromptMetadata(McpPromptMetadata... promptMetadataCollection) {
        for (McpPromptMetadata mcpPromptMetadata : promptMetadataCollection) {
            promptsMap.put(mcpPromptMetadata.getPrompt().getPrompt().getName(), mcpPromptMetadata);
        }
    }

    /**
     * 获取提示元数据
     * @param name 提示名称
     * @return 提示元数据
     */
    public McpPromptMetadata getMcpPromptMetadata(String name) {
        return promptsMap.get(name);
    }

    /**
     * 获取资源集合
     * @return 资源集合
     */
    public List<Resource> getResources() {
        return Collections.unmodifiableList(resources.stream().map(McpResourceMetadata::getResource).collect(Collectors.toList()));
    }

    /**
     * 获取资源元数据集合
     * @return 资源元数据集合
     */
    public List<McpResourceMetadata> getResourcesMetadata() {
        return Collections.unmodifiableList(resources);
    }

    /**
     * 添加资源元数据
     * @param resourceMetadataCollection 资源元数据集合
     */
    public void addResourceMetadata(McpResourceMetadata... resourceMetadataCollection) {
        for (McpResourceMetadata mcpResourceMetadata : resourceMetadataCollection) {
            ResourceSchema.BaseResource resource = mcpResourceMetadata.getResource().getResource();
            String uniqueUri = resource instanceof ResourceSchema.Resource
                    ? ((ResourceSchema.Resource) resource).getUri()
                    : ((ResourceSchema.ResourceTemplate) resource).getUriTemplate();
            if (!resoureUriSet.add(uniqueUri)) {
                throw new McpServerException("Resource Uri[" + uniqueUri + "] Already Exists.");
            }
            resources.add(mcpResourceMetadata);
        }
    }

    /**
     * 添加资源
     * @param uri 资源uri
     * @param name 资源name
     * @param description 资源描述
     * @param mimeType 资源类型
     * @param data 资源数据
     */
    public void addResource(String uri, String name, String description, String mimeType, String data) {
        McpResourceMetadata mcpResourceMetadata = createMcpResourceMetadata(uri, name, description, mimeType);
        mcpResourceMetadata.setTextContent(true);
        mcpResourceMetadata.setData(data);
        addResourceMetadata(mcpResourceMetadata);
    }

    /**
     * 添加资源
     * @param uri 资源uri
     * @param name 资源name
     * @param description 资源描述
     * @param mediaType 媒体类型
     * @param data 资源数据
     */
    public void addResource(String uri, String name, String description, MediaType mediaType, String data) {
        addResource(uri, name, description, mediaType.getValue(), data);
    }

    /**
     * 添加资源
     * @param uri 资源uri
     * @param name 资源name
     * @param description 资源描述
     * @param mimeType 资源类型
     * @param data 资源数据
     */
    public void addResource(String uri, String name, String description, String mimeType, byte[] data) {
        McpResourceMetadata mcpResourceMetadata = createMcpResourceMetadata(uri, name, description, mimeType);
        mcpResourceMetadata.setTextContent(false);
        mcpResourceMetadata.setData(new String(Base64.getEncoder().encode(data), StandardCharsets.UTF_8));
        addResourceMetadata(mcpResourceMetadata);
    }

    /**
     * 添加资源
     * @param uri 资源uri
     * @param name 资源name
     * @param description 资源描述
     * @param mediaType 媒体类型
     * @param data 资源数据
     */
    public void addResource(String uri, String name, String description, MediaType mediaType, byte[] data) {
        addResource(uri, name, description, mediaType.getValue(), data);
    }

    /**
     * 创建资源元数据
     * @param uri 资源uri
     * @param name 资源name
     * @param description 资源描述
     * @param mimeType 资源类型
     * @return 资源元数据
     */
    private McpResourceMetadata createMcpResourceMetadata(String uri, String name, String description, String mimeType) {
        ResourceSchema.Resource schemaResource = new ResourceSchema.Resource();
        schemaResource.setUri(uri);
        schemaResource.setName(name);
        schemaResource.setDescription(description);
        schemaResource.setMimeType(mimeType);
        McpResourceMetadata mcpResourceMetadata = new McpResourceMetadata();
        Resource resource = new Resource();
        resource.setResource(schemaResource);
        mcpResourceMetadata.setResource(resource);
        return mcpResourceMetadata;
    }
}
